package cn.z2huo.demo.hello.kafka.support;

import cn.z2huo.demo.hello.core.common.TransactionCodeHolder;
import cn.z2huo.demo.hello.core.common.util.TransactionUtils;
import cn.z2huo.demo.hello.core.exception.KafkaConsumeException;
import cn.z2huo.demo.hello.kafka.api.KafkaConnectionFactory;
import cn.z2huo.demo.hello.kafka.api.KafkaConsumable;
import cn.z2huo.demo.hello.kafka.api.KafkaSubscriber;
import cn.z2huo.demo.hello.kafka.model.KafkaOffset;
import cn.z2huo.demo.hello.kafka.service.api.KafkaFailService;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.springframework.context.ApplicationContext;

import java.time.Instant;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * <p>
 *
 * @author z2huo
 */
@Slf4j
public class ConsumerRunner implements KafkaConsumable {

    /**
     * 用于控制线程关闭的开关
     */
    private final AtomicBoolean closed = new AtomicBoolean(false);

    private final String consumerName;

    private KafkaConsumer<String, String> consumer;

    private final KafkaSubscriber subscriber;

    private final KafkaFailService kafkaFailService;

    private final KafkaConnectionFactory connectionFactory;

    private Long commitInterval;

    public ConsumerRunner(String consumerName, KafkaSubscriber subscriber, ApplicationContext applicationContext, KafkaConnectionFactory factory) {
        this.consumerName = consumerName;
        this.subscriber = subscriber;
        this.connectionFactory = factory;

        if (subscriber.errorContinue()) {
            this.kafkaFailService = applicationContext.getBean(KafkaFailService.class);
        } else {
            this.kafkaFailService = null;
        }
        if (subscriber.commitIntervalSize() < 1) {
            throw new KafkaConsumeException(-1, "commitInterval size must be greater than 0");
        }

    }

    @Override
    public void shutdown() {
        try {
            closed.set(true);
            consumer.wakeup();
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }
    }

    @Override
    public void run() {
        boolean restart = false;
        try {
            // 通过开关控制是否需要持续监听队列
            while (!closed.get()) {
                // 如果订阅配置要求消费者暂停获取消息，则暂停不做消息获取
                if (subscriber.needPause()) {
                    continue;
                }

                // 每次的消费设置一个交易码，可以被覆盖
                TransactionCodeHolder.set(TransactionUtils.generateTransactionCode());

                ConsumerRecords<String, String> records = consumer.poll(subscriber.timeout());
                for (TopicPartition partition : records.partitions()) {
                    // 获取分区中的记录集合
                    List<ConsumerRecord<String, String>> recordList = records.records(partition);
                    // 根据客户端设置的每次处理笔数进行循环处理
                    if (CollectionUtils.isNotEmpty(recordList)) {
                        Instant now = this.getNow();

                        // 根据 lastOffset 判断消费是否出现异常，若出现异常则跳出循环，重新拉取消费
                        long lastOffset = 0;
                        for (int i = 0; i < recordList.size(); i++) {
                            // 消费
                            lastOffset = execConsumer(partition, List.of(recordList.get(i)));
                            // 如果为 -1，则不提交 offset
                            if (lastOffset == -1) {
                                if (subscriber.errorContinue()) {
                                    // 持久化当前消息，预先强制清空动态数据源信息
//                                    DynamicDataSourceHolder.clear();
                                    if (!kafkaFailService.onFail(recordList.get(i), consumerName)) {
                                        // 消息持久化失败，需要重新维护客户端记录的 offset
                                        consumer.seek(partition, recordList.get(i).offset());
                                        break;
                                    } else {
                                        // 此次消费成功
                                        lastOffset = recordList.get(i).offset();
                                    }
                                } else {
                                    // 若内部消息出现异常，则重新维护客户端记录的 offset
                                    consumer.seek(partition, recordList.get(i).offset());
                                    break;
                                }
                            }

                            if ((now != null && now.isBefore(Instant.now())) || i % subscriber.commitIntervalSize() == 0 || i == recordList.size() - 1) {
                                KafkaOffset kafkaOffset = new KafkaOffset();
                                kafkaOffset.setConsumerName(consumerName);
                                kafkaOffset.setGroupId(subscriber.groupId());
                                kafkaOffset.setTopicPartition(partition);
                                kafkaOffset.setOffset(lastOffset + 1);
                                this.syncOffset(kafkaOffset, List.of(recordList.get(i)));
                                now = this.getNow();
                            }
                        }

                        if (lastOffset == -1) {
                            break;
                        }
                    }
                }
            }
        } catch (Exception e) {
            if (StringUtils.isNotBlank(e.getMessage()) && StringUtils.contains(e.getMessage(), "rebalance")) {
                restart = true;
            }
        } finally {
            consumer.close();
            if (restart) {
                this.restart();
            }
        }
    }

    private Instant getNow() {
        if (subscriber.commitIntervalTime() > 0) {
            return Instant.ofEpochMilli(Instant.now().toEpochMilli() + subscriber.commitIntervalTime());
        }
        return null;
    }

    private void restart() {
        try {
            this.prepareConsumer();
            this.run();
        } catch (Exception e) {
            log.error("重新加载 ConsumerRunner 失败，该节点无法订阅 kafka 消息，data = {}", consumerName);
            log.error("重新加载 ConsumerRunner 失败，该节点无法订阅 kafka 消息", e);
        }
    }

    /**
     * <p>同步 offset 内容
     * <p>如果需要做 offset 外部存储，则构造对象，调用 DAO 来存储 offset
     * <p>如果需要提交至 kafka，则获取最后一个 offset，并 +1 提交
     *
     * @param kafkaOffset   offset 信息
     * @param records   需要消费的记录数
     */
    private void syncOffset(KafkaOffset kafkaOffset, List<ConsumerRecord<String, String>> records) {
        try {
            // 如果需要提交，则获取最后一个 offset，并 +1 提交
            if (subscriber.getOffsetDAO() == null) {
                consumer.commitSync(Map.of(kafkaOffset.getTopicPartition(), new OffsetAndMetadata(kafkaOffset.getOffset())));
            } else {
                subscriber.getOffsetDAO().store(kafkaOffset);
            }
        } catch (Exception e) {
            subscriber.rollback(records, kafkaOffset);
            throw e;
        }
    }

    private void prepareConsumer() {
        Properties properties = new Properties();
        properties.put("group.id", subscriber.groupId());
        this.consumer = connectionFactory.getConsumer(consumerName, properties);

        try {
            consumer.subscribe(subscriber.topics());
            if (subscriber.getOffsetDAO() != null) {
                subscriber.getOffsetDAO().list(consumerName, subscriber.groupId(), subscriber.topics())
                        .forEach(kafkaOffset -> consumer.seek(kafkaOffset.getTopicPartition(), kafkaOffset.getOffset()));
            }
        } catch (Exception e) {
            consumer.close();
            throw e;
        }
    }

    private long execConsumer(TopicPartition topicPartition, List<ConsumerRecord<String, String>> records) {
        try {
            subscriber.consume(records, topicPartition);
            // 计算消息集合的最后一位 offset
            return records.getLast().offset();
        } catch (KafkaConsumeException e) {
            log.error("occur an error with topic = " + topicPartition.topic(), e);
            if (e.getLastConsumeRecordIndex() != -1) {
                return records.get(e.getLastConsumeRecordIndex()).offset();
            }
            return -1;
        } catch (Exception e) {
            log.error("occur an error with topic = " + topicPartition.topic(), e);
            return -1;
        }
    }

}
